#include "XYZJoint.h"
#include "Mat6x6.h"
#include "PrescribedMotion.h"
#include "RigidBodyNode.h"
#include "EOMData.h"
#include "ABAData.h"

using namespace RSIM;
using namespace std;

///////////////////////////////////////////////////////

int XYZ::getJointDOF()const{return 3;}

///////////////////////////////////////////////////////

void XYZ::calc_H_x_UDot(const State& state,double *h_x_udot) const{
	const double *UDot = state.getUDotForNode(this->NodeIndex_);
	h_x_udot[3] = UDot[0];
	h_x_udot[4] = UDot[1];
	h_x_udot[5] = UDot[2];
}

///////////////////////////////////////////////////////

void XYZ::setDefaultQ(double *q)const{
	q[0] = 0.0;	q[1] = 0.0;
	q[2] = 0.0;
}

///////////////////////////////////////////////////////

void XYZ::setDefaultU(double *u)const{
	u[0] = 0.0;	u[1] = 0.0;
	u[2] = 0.0;
}

///////////////////////////////////////////////////////

void XYZ::calcQDot(State& state)const{
	state.wgetQDotForNode(this->NodeIndex_)[0] = state.wgetUForNode(this->NodeIndex_)[0];
	state.wgetQDotForNode(this->NodeIndex_)[1] = state.wgetUForNode(this->NodeIndex_)[1];
	state.wgetQDotForNode(this->NodeIndex_)[2] = state.wgetUForNode(this->NodeIndex_)[2];
}

///////////////////////////////////////////////////////

void XYZ::updRelVelocity(const State& state,Vect3& w, Vect3& v)const{
	v(0) = state.getUForNode(this->NodeIndex_)[0];
	v(1) = state.getUForNode(this->NodeIndex_)[1];
	v(2) = state.getUForNode(this->NodeIndex_)[2];
}

///////////////////////////////////////////////////////

XYZ::XYZ(const int& NodeIndex):Joint(XYZJoint){
	this->NodeIndex_ = NodeIndex;
}

///////////////////////////////////////////////////////

XYZ* XYZ::New(const int& NodeIndex){
	return new XYZ(NodeIndex);
}

///////////////////////////////////////////////////////

void XYZ::printJointType() const{
	cout<<"Joint Type=XYZ\n";
}

///////////////////////////////////////////////////////

void XYZ::updJointTransform(	const double *q,
				const RigidBodyNode *Parent,
				const int& pChildID,
				double *betak){
	double tmp[] = {q[0],q[1],q[2]};
	this->T_MF_.wgetP().copy(tmp);
	
	this->updBetaK(q,Parent,pChildID,betak);
}

///////////////////////////////////////////////////////

void XYZ::updBetaK(	const double *q,
			const RigidBodyNode *Parent,
			const int& pChildID,
			double *betak)const{
	if(Parent){
		const double *T_FM_P_Ptr = Parent->data_->getBodyTransformFM()[pChildID].P().Ptr();
		const double *T_FM_R_Ptr = Parent->data_->getBodyTransformFM()[pChildID].R().Ptr();

		FastOps::mult_mat33_vec3(T_FM_R_Ptr,q,betak);
		
		betak[0] += T_FM_P_Ptr[0];
		betak[1] += T_FM_P_Ptr[1];
		betak[2] += T_FM_P_Ptr[2];						
	}	
}

///////////////////////////////////////////////////////

void XYZ::updVelocity(	const State& state, 
			const RigidBodyNode *Parent,
			const Rotation& prF_C_F, 
			const double *betak,
			Vect3& w,
			Vect3& v) const {
					    
	const Vect3& pr_w = Parent->data_->getAngularVelocity();
	const Vect3& pr_v = Parent->data_->getLinearVelocity();
	
	// 1. F_C_prF*pr_w is the absolute angular velocity
	// of the parent expressed in the joint's F frame
	double *wptr = w.wPtr();
	FastOps::transpose_mult_mat33_vec3(prF_C_F.Ptr(),pr_w.Ptr(),wptr);
	
	double v_tmp[3];
	
	// omega x r
	FastOps::cross(pr_w.Ptr(),betak,v_tmp);
	
	// v + omega x r
	v_tmp[0] += pr_v(0);
	v_tmp[1] += pr_v(1);
	v_tmp[2] += pr_v(2);
	
	double *v_wptr = v.wPtr();
	
	// convert v + omega x r to F basis of this body
	FastOps::transpose_mult_mat33_vec3(prF_C_F.Ptr(), v_tmp, v_wptr);
	
	// Velocity = v + omega x r + qdot
	const double *u_tmp = state.getUForNode(this->NodeIndex_);
	v_wptr[0] += u_tmp[0];
	v_wptr[1] += u_tmp[1];
	v_wptr[2] += u_tmp[2];
}

///////////////////////////////////////////////////////

void XYZ::calcAKT(	const RigidBodyNode *Parent, 
			const RigidBodyNode *Body,
			const int& ChildID,
			const double *betak,
			State& state,
			Vect6& sAkt) const{
	this->calcQDot(state);
	double *UDot = state.wgetUDotForNode(this->NodeIndex_);		
	if(Parent){
		const double *U = state.getUForNode(this->NodeIndex_);
		
		// AKT is calculated in F basis of this body
		double *akt_ptr = sAkt.wPtr();
		
		// set Akt to zero to clear existing values.
		akt_ptr[0] = 0.0;akt_ptr[1] = 0.0;akt_ptr[2] = 0.0;
		
		const double *prFCF_Ptr = Body->data_->getRotationFCF().Ptr();
		
		double s_betak[3], ang_vel[3];
		
		// Angular velocity of parent expressed in F of parent
		const double *parent_ang_ptr = Parent->data_->getAngularVelocity().Ptr();
		
		// express angular velocity and betak in F basis of this body
		FastOps::transpose_mult_mat33_vec3(prFCF_Ptr,parent_ang_ptr,ang_vel);
		FastOps::transpose_mult_mat33_vec3(prFCF_Ptr,betak,s_betak);
		
		// cross( w^{k-1}, cross(w^{k-1}, \beta^{k-1}) ) {expressed in F of this body}		
		FastOps::wxwxr(ang_vel, s_betak, akt_ptr+3);				
		
		double wxvkk[3];
		// 2*cross(w^k, v^{k+1}_{k+1}u^k+1)
		FastOps::cross(ang_vel,U,wxvkk);
		akt_ptr[3] += 2*wxvkk[0];
		akt_ptr[4] += 2*wxvkk[1];
		akt_ptr[5] += 2*wxvkk[2];
		
		// In this case, Akt = Akt + p^k_k \dot{u}^k
		if(this->isMotionPrescribedAndActive_){
			// Update UDot_ cache from prescribed motion
			this->PM_->setUDot(UDot);
			akt_ptr[3] += UDot[0];
			akt_ptr[4] += UDot[1];
			akt_ptr[5] += UDot[2];
		}
	} // if(Parent)
	else{
		if(this->isMotionPrescribedAndActive_){
			this->PM_->setUDot(UDot);
		}
	}
}

///////////////////////////////////////////////////////

void XYZ::sweepBackwards(	const bool& isLeaf, 
				const int& ChildID,
				const Vect6& sAkt,
				const Rotation& FCF,					    
				const EOMData *eom,
				const double *betak,
				RigidBodyNode *Parent,
				ABAData *aba) const{
				  
	// In the following implementation, I have stepped through intermediate frames
	// hence three calls to calc_Triang functions instead of just one. 
	// This was indeed avoidable but would require the inversion of a dofxdof (3x3 in this case)
	// matrix. I am uncertain as to which method would result in a lower operation count. 
	// Need to work this out at somepoint.
	
	// get writable pointers for articulated body inertia of this body.
	// Ik3[0] is the ABI of this body and the eqns of motion for this body
	// can be written as Ik3[0] * A - Fk3[0] = 0. Here the ABI of intermediate frames 
	// are calculated and stored in Ik3+1 and Ik3+2
	Mat6x6 *Ik3 = aba->wgetIk3();
	Vect6 *Fk3 = aba->wgetFk3();	
	
	if(!isMotionPrescribedAndActive_){
		
		this->calc_Triang_x_Ik3_TransZ(Ik3->Ptr(),(Ik3+1)->wPtr());
		this->calc_Triang_x_cFk_TransZ(Ik3->Ptr(), Fk3->Ptr(), (Fk3+1)->wPtr());		

		this->calc_Triang_x_Ik3_TransY((Ik3+1)->Ptr(), (Ik3+2)->wPtr());	
 		this->calc_Triang_x_cFk_TransY((Ik3+1)->Ptr(), (Fk3+1)->Ptr(), (Fk3+2)->wPtr());

		if(Parent){
			this->calc_cFk((Fk3+2)->Ptr(),(Ik3+2)->Ptr(),sAkt.Ptr(),aba->wgetcFk()->wPtr());			
			double triang_ik3[36],triang_fk3[6];
			this->calc_Triang_x_Ik3_TransX((Ik3+2)->Ptr(),triang_ik3);
			this->calc_Triang_x_cFk_TransX((Ik3+2)->Ptr(),aba->getcFk()->Ptr(),triang_fk3);
			
			double sIk3[36], sFk3[6];
			FastOps::spatial_basis_shift_mat(FCF.Ptr(), triang_ik3, sIk3);
			FastOps::spatial_basis_shift_vec(FCF.Ptr(), triang_fk3, sFk3);
			

			double Ik3Child[36], Fk3Child[6];
			this->calc_skbeta_x_Ik3_x_tskbeta(sIk3, betak, Ik3Child);
			this->calc_skbeta_x_cFk(sFk3, betak, Fk3Child);

			Parent->data_->wgetABAData()->wgetIk3()[0].add(Ik3Child);
			Parent->data_->wgetABAData()->wgetFk3()[0].add(Fk3Child);
			
		}// If (Parent)
	} // if (!isMotionPrescribedAndActive)
	else{
		if(Parent){
			// cFk = Fk3 - Ik3*A^{k-1}_t
			this->calc_cFk(Fk3->Ptr(),Ik3->Ptr(),sAkt.Ptr(),aba->wgetcFk()->wPtr());
			
			double sIk3[36], sFk3[6];
			FastOps::spatial_basis_shift_mat(FCF.Ptr(), Ik3->Ptr(), sIk3);
			FastOps::spatial_basis_shift_vec(FCF.Ptr(), aba->getcFk()->Ptr(), sFk3);
			
			double Ik3Child[36], Fk3Child[6];
			this->calc_skbeta_x_Ik3_x_tskbeta(sIk3, betak, Ik3Child);
			this->calc_skbeta_x_cFk(sFk3, betak, Fk3Child);
			
			Parent->data_->wgetABAData()->wgetIk3()[0].add(Ik3Child);
			Parent->data_->wgetABAData()->wgetFk3()[0].add(Fk3Child);
		}
	}
}

///////////////////////////////////////////////////////

void XYZ::sweepForward( const RigidBodyNode *Parent,
			const int& ChildID,
			const Vect6& sAkt,
			const Rotation& FCF,
			const ABAData *aba,
			const double *betak,
			EOMData *eom,
			State& state
			){				
	double *UDot = state.wgetUDotForNode(this->NodeIndex_);	
	
	const Vect6 *Fk3 = aba->getFk3();
	const Mat6x6 *Ik3 = aba->getIk3();
	Vect6 *Ak1 = eom->wgetAk1();
	double *ak1_ptr = Ak1->wPtr();	

	if(Parent){
		const double *Parent_Ak1_Ptr = Parent->data_->getEOMData()->getAk1()->Ptr();
		if(!isMotionPrescribedAndActive_){
			{
				double ak1_tmp[6];
				FastOps::transpose_skbeta_x_V(Parent_Ak1_Ptr, betak, ak1_tmp);
				FastOps::transpose_spatial_basis_shift_vec(FCF.Ptr(),ak1_tmp, ak1_ptr);
				this->calc_UDot_TransX((Ik3+2)->Ptr(),ak1_ptr, aba->getcFk()->Ptr(), UDot);
			}			
			ak1_ptr[3] += UDot[0];
			Ak1->add(sAkt.Ptr());			
			
			this->calc_UDot_TransY((Ik3+1)->Ptr(),ak1_ptr,(Fk3+1)->Ptr(),UDot+1);
			ak1_ptr[4] += UDot[1];

			this->calc_UDot_TransZ(Ik3->Ptr(),ak1_ptr,Fk3->Ptr(),UDot+2);
			ak1_ptr[5] += UDot[2];
		}//if(!isMotionPrescribedAndActive_)
		else{
			double Ak1p1[6];
			// Ak1p1 (0:5)= transpose(skbeta)*A^{parent}_1 
			// In this case, sAkt already includes p^k_k \dot{u}^k
			FastOps::transpose_skbeta_x_V(Parent_Ak1_Ptr, betak,Ak1p1);
			
			// Ak1 = Ak1p1 (0:5) exprssed in F of this body
			FastOps::transpose_spatial_basis_shift_vec(FCF.Ptr(),Ak1p1,eom->wgetAk1()->wPtr());
			eom->wgetAk1()->add(sAkt.Ptr());
		}
	}//if(Parent)
	else{
		if(!isMotionPrescribedAndActive_){			
			const double *fk3_ptr = Fk3[2].Ptr();
			UDot[0] = fk3_ptr[3]/Ik3[2].Ptr()[21];

			double ak1_tmp[]={0.0,0.0,0.0,UDot[0],0.0,0.0};
			Ak1->copy(ak1_tmp);
			
			this->calc_UDot_TransY((Ik3+1)->Ptr(),ak1_ptr,(Fk3+1)->Ptr(),UDot+1);
			ak1_ptr[4] = UDot[1];
						
			this->calc_UDot_TransZ(Ik3->Ptr(),ak1_ptr,Fk3->Ptr(),UDot+2);			
			ak1_ptr[5] = UDot[2];			
		}
		else{
			// Calculate Ak1_ in F frame of this body
			double *ak1_ptr = eom->wgetAk1()->wPtr();			
			ak1_ptr[1]=0.0; ak1_ptr[2]=0.0; ak1_ptr[3]=0.0;
			ak1_ptr[3]=UDot[0]; ak1_ptr[4]=UDot[1]; ak1_ptr[5]=UDot[2];
		}		
	}
	
	this->calc_m_cFk(aba->getFk3()[0].Ptr(), aba->getIk3()[0].Ptr(), Ak1->Ptr(), eom->wgetFkc1()->wPtr());	
}

///////////////////////////////////////////////////////